#ifndef TESTGRAMMAR_LISTENER_H
#define TESTGRAMMAR_LISTENER_H
#include "TestGrammarParserBaseListener.h"
#include "AST/UniASTNode.h"
#include <any>
#include <stack>

using namespace UniAST;
class TestGrammarListener : public TestGrammarParserBaseListener {

private:
  std::map<antlr4::ParserRuleContext *, std::string> &astmap;
  RuleNode *root;
  std::stack<TreeNode *> nodeManager;

  RuleNode *_buildConstructTreeNode(
      std::pair<std::string, std::vector<std::pair<std::string, std::any>>>
          extraDecl) {
    auto extraDeclName = extraDecl.first;
    auto extraNode = new RuleNode(nullptr, extraDeclName);
    auto thisNode = this->nodeManager.top();
    for (auto &extraArg : extraDecl.second) {
      auto argType = extraArg.first;
      if (argType == "NODES") {
        auto args = std::any_cast<std::vector<antlr4::ParserRuleContext *>>(
            extraArg.second);
        for (auto ctx : args) {
          auto child = thisNode->findChild(ctx);
          thisNode->deleteChild(child);
          extraNode->addChild(child);
        }
        thisNode->addChild(extraNode);
      } else if (argType == "STRING_LITERAL") {
        std::cout << "TODO: argType == STRING_LITERAL " << std::endl;
      } else if (argType == "TYPE") {
        std::cout << "TODO: argType == TYPE " << std::endl;
      } else if (argType == "CONSTRUCTOR") {
        auto arg = std::any_cast<std::pair<
            std::string, std::vector<std::pair<std::string, std::any>>>>(
            extraArg.second);
        auto child = this->_buildConstructTreeNode(arg);
        thisNode->deleteChild(child);
        extraNode->addChild(child);
        thisNode->addChild(extraNode);
      } else {
        std::cerr << "unknown arg type: " << argType << std::endl;
        exit(1);
      }
    }
    return extraNode;
  }

public:

	TestGrammarListener(std::map<antlr4::ParserRuleContext *, std::string> &astmap)
		: astmap(astmap), root(new RuleNode(nullptr, "root")) {
		this->nodeManager.push(this->root);
	};

	~TestGrammarListener() { 
		delete root; 
	}


  RuleNode *getRoot() { 
    return this->root; 
  }

  void printStringTree() {
    std::cout << this->root->toStringTree() << std::endl;
  }

  virtual void enterEveryRule(antlr4::ParserRuleContext *ctx) override {
    if (!this->astmap.contains(ctx)) {
      return;
    }
    auto nodeName = this->astmap[ctx];
    auto newNode = new RuleNode(ctx, nodeName);
    auto p = this->nodeManager.top();
    p->addChild(newNode);
    this->nodeManager.push(newNode);
  }
  virtual void exitEveryRule(antlr4::ParserRuleContext *ctx) override {
    if (!this->astmap.contains(ctx)) {
      return;
    }
    this->nodeManager.pop();
  }


	virtual void exitAdd_expr(TestGrammarParser::Add_exprContext *ctx) override {
		for (auto construct : ctx->extra_node_construct) {
			this->_buildConstructTreeNode(construct);
		}
	}

};

#endif